"""
python fujian1_AutoTsArima_after_smooth.py
"""

import os
import json
import pandas as pd
from autots import AutoTS

# 设置输入和输出目录（使用相对路径）
input_dir = r"./fujian/fujian1/spline_then_smooth"
output_dir = r"./fujian/fujian1/smooth_then_AutoTsArima"

# 创建输出目录（如果不存在）
os.makedirs(output_dir, exist_ok=True)

# 遍历目录中的所有 JSON 文件
for filename in os.listdir(input_dir):
    if filename.endswith('.json'):
        file_path = os.path.join(input_dir, filename)

        # 读取 JSON 文件
        with open(file_path, 'r') as f:
            data = json.load(f)

        # 将数据转换为 DataFrame
        df = pd.DataFrame(data)
        df['date'] = pd.to_datetime(df['date'])  # 转换日期格式
        df.set_index('date', inplace=True)

        # 确保有 12 个数据点进行训练
        if len(df) >= 12:
            # 使用 AutoTS 自动选择 ARIMA 参数并进行预测
            model = AutoTS(forecast_length=3, frequency='M', ensemble='simple', models='ARIMA')
            model.fit(df.reset_index(), date_col='date', value_col='inventory')

            # 进行预测
            forecast = model.predict()
            forecast_df = forecast['predictions']

            # 生成未来日期
            forecast_dates = pd.date_range(start='2023-07-01', periods=3, freq='M')

            # 创建预测结果的 DataFrame
            forecast_result = pd.DataFrame({'date': forecast_dates, 'inventory': forecast_df.values.flatten()})

            # 合并原始数据和预测数据
            result_df = pd.concat([df.reset_index(), forecast_result], ignore_index=True)

            # 转换回 JSON 格式并输出
            output_path = os.path.join(output_dir, filename)
            result_json = result_df.to_dict(orient='records')

            with open(output_path, 'w') as f:
                json.dump(result_json, f, indent=4)

            print(f"Processed {filename} and saved results to {output_path}")
        else:
            print(f"Not enough data points in {filename} for ARIMA model.")
